/*
* This file is a part of the open-eBackup project.
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
* If a copy of the MPL was not distributed with this file, You can obtain one at
* http://mozilla.org/MPL/2.0/.
*
* Copyright (c) [2024] Huawei Technologies Co.,Ltd.
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
*/
#include "VssRestore.h"

#include <map>
#include <memory.h>
#include <cstdio>
#include <memory>
#include <comutil.h>
#include <iostream>
#include <fstream>
#include <atlstr.h>


namespace {
    const int POSTFIX_POS = 4;
    const int WILDCARD_POS = 5;
}

VssRestore::VssRestore()
{
}

VssRestore::~VssRestore()
{
}

bool VssRestore::Restore(const std::string& backupXmlPath, const std::string& writerXmlPath, const std::string& guid, \
    const std::string& dataRepoPath, const std::string& restorePath)
{
    // Backup Metadata Document
    std::wstring xmlContent;
    DWORD retd = LoadFromFile(backupXmlPath, xmlContent);
    if (retd != 0) {
        return false;
    }
    CComBSTR Bstr(xmlContent.size(), xmlContent.data());
    m_backupMetadataDocument = Bstr;

    std::wstring xmlContentWriter;
    // Exchange Write Metadata Document
    retd = LoadFromFile(writerXmlPath, xmlContentWriter);
    if (retd != 0) {
        return false;
    }
    CComBSTR BstrWriter(xmlContentWriter.size(), xmlContentWriter.data());
    m_exchangeWriterMetadataDocument = BstrWriter;

    // VSS流程
    m_targetDatabaseGuid = _com_util::ConvertStringToBSTR(guid.c_str());
    VSS_ID exchangeWriterID;
    bool ret = RestoreInitalization(m_backupMetadataDocument);
    if (!ret) {
        return false;
    }
    ret = RestorePreparetionV2(exchangeWriterID, m_targetDatabaseGuid);
    if (!ret) {
        return false;
    }
    ret = ActualRestoration(dataRepoPath, restorePath);
    if (!ret) {
        return false;
    }
    ret = RestoreCleanUpAndTermination();
    if (!ret) {
        return false;
    }

    return true;
}

bool VssRestore::RestoreInitalization(BSTR backupMetadataDocument)
{
    // 初始化
    HRESULT hr = ::CoInitializeEx(nullptr, COINIT::COINIT_MULTITHREADED);
    CHECK_HR_RETURN(hr, "CoInitializeEx", false);

    hr = CoInitializeSecurity(nullptr,  //  Allow *all* VSS writers to communicate back!
        -1,                             //  Default COM authentication service
        nullptr,                        //  Default COM authorization service
        nullptr,                        //  reserved parameter
        RPC_C_AUTHN_LEVEL_PKT_PRIVACY,  //  Strongest COM authentication level
        RPC_C_IMP_LEVEL_IMPERSONATE,    //  Minimal impersonation abilities
        nullptr,                        //  Default COM authentication settings
        EOAC_DYNAMIC_CLOAKING,          //  Cloaking
        nullptr                         //  Reserved parameter
    );
    CHECK_HR_RETURN(hr, "CoInitializeSecurity", false);

    // 创建实例
    hr = S_OK;
    hr = CreateVssBackupComponents(&m_ppBackupComponent);
    CHECK_HR_RETURN(hr, "CreateVssBackupComponents", false);

    // 装载备份阶段生成的XML
    hr = m_ppBackupComponent->InitializeForRestore(m_backupMetadataDocument);
    CHECK_HR_RETURN(hr, "InitializeForRestore", false);

    // 创建writer实例
    hr = S_OK;
    hr = CreateVssExamineWriterMetadata(m_exchangeWriterMetadataDocument, &m_ppExchangeWriterComponentPlus);
    CHECK_HR_RETURN(hr, "CreateVssExamineWriterMetadata", false);

    // 通知各writers将各自的metadata发送过来
    CComPtr<IVssAsync> pAsync;
    hr = m_ppBackupComponent->GatherWriterMetadata(&pAsync);
    CHECK_HR_RETURN(hr, "GatherWriterMetadata", false);

    CHECK_BOOL_RETURN(WaitAndCheckForAsyncOperation(pAsync), "Wait GatherWriterMetadata", false);

    // 获取备份阶段使用的writer
    hr = m_ppBackupComponent->GetWriterComponentsCount(&m_pcWriterComponents);
    CHECK_HR_RETURN(hr, "GetWriterComponentsCount", false);

    for (UINT i = 0; i < m_pcWriterComponents; ++i) {
        // 保存所有使用的writer
        VSS_ID pidInstance;
        CComPtr<IVssExamineWriterMetadata> ppExamineWriterMetadata;
        hr = m_ppBackupComponent->GetWriterMetadata(i, &pidInstance, &ppExamineWriterMetadata);
        CHECK_HR_RETURN(hr, "GetWriterMetadata", false);
        m_pidWriterInstances.push_back(pidInstance);
        m_ppExmaineWriterMetadata.push_back(ppExamineWriterMetadata);
    }
    
    return true;
}

// dataDir: 包含通配符的绝对路径
// fileList: 所有匹配文件的文件名(而不是绝对路径)
// 扫描dataDir所在文件夹, 将所有匹配文件的文件名放入fileList中返回
void VssRestore::GetFileList(std::string dataDir, std::vector<std::wstring> &fileList)
{
    HANDLE hFind;
    WIN32_FIND_DATA data;
    hFind = FindFirstFileA(dataDir.c_str(), &data);
    if (hFind != INVALID_HANDLE_VALUE) {
        do {
            std::string tempName = data.cFileName;
            if (tempName == "." || tempName == ".." || (data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)) {
            } else {
                std::wstring wTempName = Utf8ToUtf16(tempName);
                fileList.emplace_back(wTempName);
            }
        } while (FindNextFileA(hFind, &data));
        FindClose(hFind);
    } else {
        auto errCode = GetLastError();
    }
}

bool VssRestore::ActualRestoration(const std::string& dataRepoPath, const std::string& restorePath)
{
    // 生成Restore Set
    std::map<std::string, std::string> restorePathList;     // (src, des)

    // 注意传入路径结尾是否有反斜杠, 这里暂时默认有
    for (int i = 0; i < m_bstrFilenameToRestore.size(); ++i) {
        std::string fileNameStr = _com_util::ConvertBSTRToString(m_bstrFilenameToRestore[i]);
        std::string filePostfix = fileNameStr.substr(fileNameStr.size() - POSTFIX_POS, POSTFIX_POS);
        std::string srcPath = dataRepoPath + fileNameStr;
        std::string desPath = restorePath + fileNameStr;
        
        if (filePostfix == ".log" && fileNameStr[fileNameStr.size() - WILDCARD_POS] == '*') {
            // 日志文件
            // /path/to/data/repo + *.log
            std::vector<std::wstring> fileList;
            GetFileList(srcPath, fileList);
            for (int j = 0; j < fileList.size(); ++j) {
                std::string realSrcPath = dataRepoPath + Utf16ToUtf8(fileList[j]);
                std::string realDesPath = restorePath + Utf16ToUtf8(fileList[j]);
                restorePathList.insert(make_pair(realSrcPath, realDesPath));
            }
        } else {
            // 数据库文件&检查点文件
            // /path/to/data/repo + .edb/.chk
            restorePathList.insert(make_pair(srcPath, desPath));
        }
    }

    // CopyFile
    int count = 0;
    for (auto &item : restorePathList) {
        bool CopyFlag = CopyFile(item.first.c_str(), item.second.c_str(), FALSE);
        if (!CopyFlag) {
            return false;
        }
        count++;
    }

    return true;
}

bool VssRestore::RestoreCleanUpAndTermination()
{
    // 触发PostRestore事件
    HRESULT hr;
    CComPtr<IVssAsync> pAsync;
    hr = m_ppBackupComponent->PostRestore(&pAsync);
    CHECK_HR_RETURN(hr, "PostRestore", false);
    CHECK_BOOL_RETURN(WaitAndCheckForAsyncOperation(pAsync), "Wait GatherWriterMetadata", false);
    return true;
}

bool VssRestore::FetchWriterMetadata(VSS_ID exchangeWriterID)
{
    HRESULT hr;
    for (UINT i = 0; i < m_pcWriterComponents; ++i) {
        CComPtr<IVssWriterComponentsExt> ppWriter;
        hr = m_ppBackupComponent->GetWriterComponents(i, &ppWriter);
        CHECK_HR_RETURN(hr, "GetWriterComponents", false);
        // 获取Exchange writer
        VSS_ID pidInstance;
        VSS_ID pidWriter;
        hr = ppWriter->GetWriterInfo(&pidInstance, &pidWriter);
        // 正常情况下只有一个writer即exchange writer
        // 但后续应加入判断机制, 防止异常情况
        m_ppExchangeExamineWriterMetadata = m_ppExmaineWriterMetadata[i];
        m_ppExchangeWriterComponentsExt = ppWriter;
        m_pidExchangeInstance = pidInstance;
        m_pidExchangeWriter = pidWriter;
    }
    return true;
}

bool VssRestore::RestorePreparetionV2(VSS_ID exchangeWriterID, BSTR targetDatabaseGuid)
{
    // 取回从Backup文档中显式包含的writer components信息
    CHECK_BOOL_RETURN(FetchWriterMetadata(exchangeWriterID), "FetchWriterMetadata", false);

    // 遍历已加载的writer, 取得要恢复的数据库
    HRESULT hr = S_OK;
    UINT pcExchangeWriterIncludeFiles;
    UINT pcExchangeExcludeFiles;
    UINT pcExchangeComponents;
    BSTR bstrTargetComponentLogicPath;
    hr = m_ppExchangeWriterComponentPlus->GetFileCounts(&pcExchangeWriterIncludeFiles,
        &pcExchangeExcludeFiles, &pcExchangeComponents);
    CHECK_HR_RETURN(hr, "GetFileCounts", false);

    BSTR bstrExchangeRestoreComponentLogicPath;
    BSTR bstrExchangeRestoreComponentName;
    // 遍历所有component, 找到所要恢复的component
    for (UINT i = 0; i < pcExchangeComponents; ++i) {
        CComPtr<IVssWMComponent> ppComponent;
        hr = m_ppExchangeWriterComponentPlus->GetComponent(i, &ppComponent);
        CHECK_HR_RETURN(hr, "GetComponent", false);
        PVSSCOMPONENTINFO ppComponentInfo;
        hr = ppComponent->GetComponentInfo(&ppComponentInfo);
        CHECK_HR_RETURN(hr, "GetComponentInfo", false);
        // 找到所要恢复的component, 拼接其logicPath和componentName, 所得字符串即真正要恢复的component的logicPath
        std::string tempBstrComponentName = _com_util::ConvertBSTRToString(ppComponentInfo->bstrComponentName);
        std::string tempTargetDatabaseGuid = _com_util::ConvertBSTRToString(targetDatabaseGuid);
        if (tempBstrComponentName == tempTargetDatabaseGuid) {
            // 选择恢复该component
            hr = m_ppBackupComponent->SetSelectedForRestore(m_pidExchangeWriter, VSS_CT_FILEGROUP, \
                ppComponentInfo->bstrLogicalPath, ppComponentInfo->bstrComponentName, true);
            CHECK_HR_RETURN(hr, "SetSelectedForRestore", false);
            bstrExchangeRestoreComponentLogicPath = ppComponentInfo->bstrLogicalPath;
            bstrExchangeRestoreComponentName = ppComponentInfo->bstrComponentName;

            // 拼接字符串, 获取subcomponent的logicPath
            _bstr_t tempBstr1 = ppComponentInfo->bstrLogicalPath;
            _bstr_t tempBstr2 = ppComponentInfo->bstrComponentName;

            std::string tempComponentName = \
                static_cast<std::string>(_com_util::ConvertBSTRToString(ppComponentInfo->bstrComponentName));
            std::string tempLogicalPath = \
                static_cast<std::string>(_com_util::ConvertBSTRToString(ppComponentInfo->bstrLogicalPath));
            std::string finalLogicalPath = tempLogicalPath + "\\" + tempComponentName;
            bstrTargetComponentLogicPath = _com_util::ConvertStringToBSTR(finalLogicalPath.c_str());
            break;      // 若有多个writer, 不应break, 而是都放到某处记录
        }
    }

    // 再次遍历component, 找到所要恢复的数据库文件/日志文件/检查点文件
    for (UINT i = 0; i < pcExchangeComponents; ++i) {
        CComPtr<IVssWMComponent> ppComponent;
        hr = m_ppExchangeWriterComponentPlus->GetComponent(i, &ppComponent);
        PVSSCOMPONENTINFO ppComponentInfo;
        hr = ppComponent->GetComponentInfo(&ppComponentInfo);
        std::string tempBstrLogicPath = _com_util::ConvertBSTRToString(ppComponentInfo->bstrLogicalPath);
        std::string tempBstrTargetComponentLogicPath = _com_util::ConvertBSTRToString(bstrTargetComponentLogicPath);
        // 找到隐式包含的component
        if (tempBstrLogicPath == tempBstrTargetComponentLogicPath) {
            // 选择恢复所有隐式包含的component
            m_ppBackupComponent->AddRestoreSubcomponent(m_pidExchangeWriter, VSS_CT_FILEGROUP, \
                bstrExchangeRestoreComponentLogicPath, bstrExchangeRestoreComponentName, \
                ppComponentInfo->bstrLogicalPath, ppComponentInfo->bstrComponentName, false);

            // 获取待恢复文件的实际位置
            CComPtr<IVssWMFiledesc> ppFiledesc;
            BSTR pbstrPath;
            BSTR pbstrFilespec;
            for (UINT j = 0; j < ppComponentInfo->cFileCount; ++j) {
                hr = ppComponent->GetFile(j, &ppFiledesc);
                CHECK_HR_RETURN(hr, "GetFile", false);
                hr = ppFiledesc->GetPath(&pbstrPath);
                CHECK_HR_RETURN(hr, "GetPath", false);
                m_bstrPathToRestore.push_back(pbstrPath);

                hr = ppFiledesc->GetFilespec(&pbstrFilespec);
                m_bstrFilenameToRestore.push_back(pbstrFilespec);
            }
        }
    }

    // 触发PreRestore事件
    CComPtr<IVssAsync> pAsync;
    hr = m_ppBackupComponent->PreRestore(&pAsync);
    CHECK_HR_RETURN(hr, "PreRestore", false);
    CHECK_BOOL_RETURN(WaitAndCheckForAsyncOperation(pAsync), "Wait GatherWriterMetadata", false);
    return true;
}